
import torch
import os, glob
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from torch.utils.data import DataLoader
import random
import cv2.cv2 as cv2
import numpy as np


def rotate(imgPath, angle):
    # 此段代码来自百度，使用PIL旋转图片
    pilim = Image.open(imgPath)
    im2 = pilim.convert('RGBA')
    primerw, primerh = im2.size
    rot = im2.rotate(angle, expand=True)
    fff = Image.new('RGBA', rot.size, (255,) * 4)
    # 使用alpha层的rot作为掩码创建一个复合图像
    out = Image.composite(rot, fff, rot)
    out.convert(pilim.mode)
    w, h = out.size
    box = ((w - primerw) / 2, (h - primerh) / 2, (w - primerw) / 2 + primerw, (h - primerh) / 2 + primerh)
    out = out.crop(box)
    out = out.convert("RGB")
    # print(out.size)

    return out


class MyDataset3(Dataset):
    def __init__(self, imgPath, labelPath):

        self.imgPath = imgPath
        self.labelPath = labelPath

        # imgsPathList: 储存所有图片的路径
        self.imgsPathList = glob.glob(imgPath + "*.png")
        # labelsPathList: 储存所有图片对应标签的文件路径
        self.labelsPathList = [labelPath + i.split("\\")[-1][:-3] + "txt" for i in self.imgsPathList]


    def __len__(self):
        return self.imgsPathList.__len__()

    def __getitem__(self, index):
        # 获取图片路径和图片角度
        imgPath, labelPath = self.imgsPathList[index], self.labelsPathList[index]
        with open(labelPath, "r") as f:
            primerLabel = int(f.read())

        # 随机生成一个角度，当做标签，并且计算出需要对图片旋转的角度(随机角度 - 原始角度)%360
        randomAngle = random.randint(0, 359)
        label = randomAngle

        randomAngle = randomAngle - primerLabel
        randomAngle = randomAngle % 360

        tf = transforms.Compose([
            lambda x: rotate(x, randomAngle),
            transforms.Resize((320, 320)),
            transforms.ToTensor(),
            # transforms.Normalize(mean=[0.485, 0.456, 0.406],
            #                     std=[0.229, 0.224, 0.225])
        ])

        img = tf(imgPath)

        return img, label


if __name__ == '__main__':

    # 数据集测试，可以看到每个epoch都能抽出角度不一样的数据
    batchsz = 1
    dataSet = MyDataset3("data\\all_img\\", "data\\all_label\\")
    print(len(dataSet))
    dataLoader = DataLoader(dataSet, batch_size=batchsz, shuffle=True)

    for step, (x, y) in enumerate(dataLoader):
        x = torch.reshape(x, (x.shape[1], x.shape[2], x.shape[3]))
        print(x.shape)
        x = x.numpy()
        x = x.transpose(1, 2, 0)
        x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)

        print(y)
        cv2.imshow("", x)
        cv2.waitKey()
